# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: hgpark, peter

import csv
import gc
import os
import sys
import datetime
import collections
import glob
import functools

import pandas
import bisect
import matplotlib
import matplotlib.pyplot as plt
import numpy
import numexpr as ne
import json
import warnings

from difflib import SequenceMatcher
from concurrent.futures import as_completed, ProcessPoolExecutor
import coin.strategy.mm.tool.archive_base as abase
from absl import app, flags

matplotlib.use("Agg")
matplotlib.rcParams['agg.path.chunksize'] = 10000


def get_stats_list(csv_path):
  rows = []
  with open(csv_path) as f:
    csv_reader = csv.reader(f)
    header = next(csv_reader)
    for row in csv_reader:
      rows.append(row)
  return header, rows

def axvlines(xs, *args, **kwargs):
  for x in xs:
    plt.axvline(x, *args, **kwargs)


def lcs(str_a, str_b):
  lmatch = SequenceMatcher(None, str_a, str_b).find_longest_match(0, len(str_a), 0, len(str_b))
  return str_a[lmatch.a:(lmatch.a + lmatch.size)]


# Only used for plotting rmm price
def main(_):
  csv_paths = list(sorted(glob.glob(flags.FLAGS.csv_path)))
  if len(csv_paths) == 0 and flags.FLAGS.csv_path.find("csv") >= 0:
    csv_path_chunks = flags.FLAGS.csv_path.split("/")
    csv_path_chunks = csv_path_chunks[:-1] + ['*'] + csv_path_chunks[-1:]
    csv_path_str = "/".join(csv_path_chunks)
    print(csv_path_str)
    csv_paths = list(sorted(glob.glob(csv_path_str)))
  else:
    print(flags.FLAGS.csv_path)

  futures = []
  sstats = []

  compress = len(csv_paths) >= 2
  sim_elapsed_time_sum = 0
  if len(csv_paths) == 1 or flags.FLAGS.cpu == 1:
    for csv_path in csv_paths:
      header, rows = get_stats_list(csv_path)
      sstats.extend(rows)
  else:
    futures = {}
    with ProcessPoolExecutor(max_workers=min(flags.FLAGS.cpu, flags.FLAGS.max_cpu)) as executor:
      csv_path_iter = iter(csv_paths)
      while True:
        for csv_path in csv_path_iter:
          future = executor.submit(get_stats_list, csv_path)
          futures[future] = True
          if len(futures) >= flags.FLAGS.cpu:
            break
        if len(futures) == 0:
          break

        for future in as_completed(futures):
          header, rows = future.result()
          sstats.extend(rows)
          del futures[future]
          del future
          gc.collect()
          break

  sstats_df = pandas.DataFrame(sstats, columns=header)
  sstats_df['ts'] = sstats_df['ts'].astype('int64')
  sstats_df['price'] = sstats_df['price'].astype('float64')
  sstats_df = sstats_df.sort_values('ts')
  os.makedirs(flags.FLAGS.sim_result_dir, exist_ok=True)

  xlim = pandas.DatetimeIndex([sstats_df['ts'].iloc[0],sstats_df['ts'].iloc[-1]])

  def plot_graph(sim_result_dir,
               sim_result_prefix,
               strat_prefix,
               sstats_df,
               dump_stat_key,
               xlim=None,
               print_df=False,
               sim_elapsed_time=None):
    os.makedirs(sim_result_dir, exist_ok=True)
    if xlim is not None:
      plt.xlim(xlim[0], xlim[1])

    plt.rcParams["figure.figsize"] = (20, 12)

    _, ax = plt.subplots()
    for i, (label, df) in enumerate(sstats_df.groupby(['mea', 'is_midp'])):
      mid_or_index = 'mid' if label[1] == '1' else 'index'
      color = f"C{i}"
      ax = df.plot(ax=ax, kind='line', x='ts', y='price', linewidth=2, c=color, label=(label[0], mid_or_index))
      ax.set_xticklabels([datetime.datetime.fromtimestamp(ts / 1e9).strftime("%m/%d %H:%M") for ts in ax.get_xticks()])

    strat_lcs = lcs(sim_result_prefix, strat_prefix)
    if len(strat_lcs) > 10:
      short_prefix = sim_result_prefix.replace(strat_lcs, "")
    else:
      short_prefix = sim_result_prefix
    plt.suptitle(f"{strat_prefix}\n{short_prefix[:100]}\n{short_prefix[100:]}", fontsize=12)

    filename = f"s_{strat_prefix}_{short_prefix}"
    with open(f"{sim_result_dir}/filenames.txt", 'a') as txt:
      txt.write(f"{filename}\n")
    if len(filename) > 200:
      filename = filename[:200]
    filename = f"{filename}.png"
    if flags.FLAGS.sim_info:
      sim_info = json.load(open(flags.FLAGS.sim_info))
      os.makedirs(f"{sim_result_dir}/statsfile", exist_ok=True)
      with open(f"{sim_result_dir}/statsfile/{filename}.json", 'w') as sw:
        json.dump(sim_info, sw, indent=2)
    plt.savefig(f"{sim_result_dir}/" + filename)
    plt.close()

  plot_graph(flags.FLAGS.sim_result_dir,
           flags.FLAGS.sim_result_prefix,
           flags.FLAGS.strat_prefix,
           sstats_df,
           dump_stat_key="total",
           xlim=xlim,
           print_df=True,
           sim_elapsed_time=sim_elapsed_time_sum)


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  ne.set_num_threads(ne.detect_number_of_cores())

  warnings.filterwarnings("ignore", category=UserWarning)
  warnings.simplefilter(action='ignore', category=FutureWarning)

  flags.DEFINE_integer('cpu', 4, '')

  flags.DEFINE_integer('max_cpu', 4, '')

  flags.DEFINE_string('csv_path', None, '')

  flags.DEFINE_string('sim_info', None, '')

  flags.DEFINE_string('sim_result_dir', 'sim_test2', '')

  flags.DEFINE_string('sim_result_prefix', '', '')

  flags.DEFINE_string('strat_prefix', '', '')

  flags.DEFINE_bool('unfold', False, '')

  flags.DEFINE_float('split_hours', 6, '')

  plt.style.use('fivethirtyeight')
  app.run(main)
